#include "bvh_accelerator.h"

std::vector<float> as_float_array(const bvh_accelerator_t* bvh)
{
	std::vector<float> bvh_float_array;
	float pad = 0.f;
	bvh_float_array.push_back(as_float(bvh->linear_node_num));
	bvh_float_array.push_back(pad);
	bvh_float_array.push_back(pad);
	bvh_float_array.push_back(pad);

	static_assert(sizeof(linear_bvh_node_t) % (4*sizeof(float)) == 0,"sizeof linear_bvh_node_t is not aligned");
	static const int array_size = sizeof(linear_bvh_node_t)/sizeof(float);
	float buff[array_size];
	for(unsigned i = 0; i < bvh->linear_node_num; ++i)
	{
		::memcpy(buff,bvh->nodes + i,sizeof(linear_bvh_node_t));
		for(unsigned j = 0; j < array_size; ++j)
			bvh_float_array.push_back(buff[j]);
	}
	assert(bvh_float_array.size()%4 == 0);
	return bvh_float_array;
}


struct bvh_primitive_info_t
{
	bvh_primitive_info_t(int pi,const bbox_t& b)
	{
		primitive_index =pi;
		bound = b;
		
		point3f_t p;
		vadd(p,b.pmin,b.pmax);
		vsmul(centroid, 0.5f, p);
	}
	point3f_t centroid;
	bbox_t bound;
	int primitive_index;
};
struct bvh_build_node_t
{
	bvh_build_node_t(){children[0] = children[1] = NULL;bbox_init(&bounds);}
	bbox_t bounds;
	bvh_build_node_t *children[2];
	unsigned int split_axis,first_primitive_offset, n_primitives;
};

void bvh_build_node_init_leaf(bvh_build_node_t *node, unsigned offset, unsigned number_of_primitives, const bbox_t& bbox)
{
	node->first_primitive_offset = offset;
	node->n_primitives = number_of_primitives;

	node->bounds = bbox;

	node->bounds.pmax.x += 0.0001f;
	node->bounds.pmax.y += 0.0001f;
	node->bounds.pmax.z += 0.0001f;

	node->bounds.pmin.x -= 0.0001f;
	node->bounds.pmin.y -= 0.0001f;
	node->bounds.pmin.z -= 0.0001f;
}
void bvh_build_node_init_interior(bvh_build_node_t *node, unsigned dim,bvh_build_node_t* left, bvh_build_node_t* right)
{
	node->children[0] = left;
	node->children[1] = right;

	node->bounds = left->bounds;
	bbox_union_with_bbox(&node->bounds, &right->bounds);
 
	node->split_axis = dim;
	node->n_primitives = 0;

	bbox_t& bbox = node->bounds;
	bbox.pmax.x += 0.00001f;
	bbox.pmax.y += 0.00001f;
	bbox.pmax.z += 0.00001f;

	bbox.pmin.x -= 0.00001f;
	bbox.pmin.y -= 0.00001f;
	bbox.pmin.z -= 0.0001f;
}
void load_triangle(triangle_t* triangle,const shape_info_t& shape_info, const std::vector<float>& shape_data)
{
	load_triangle_points(&shape_data[0],shape_info,triangle);
}

bbox_t world_bound(const shape_info_t& shape_info, const std::vector<float>& shape_data)
{
	triangle_t triangle;
	load_triangle(&triangle,shape_info,shape_data);
	bbox_t box;
	bbox_init(&box);
	bbox_union_with_point(&box,&triangle.p0);
	bbox_union_with_point(&box,&triangle.p1);
	bbox_union_with_point(&box,&triangle.p2);
	return box;
}

bvh_build_node_t *recursive_build_bvh(bvh_accelerator_t *bvh, std::vector<bvh_primitive_info_t>& build_data,unsigned start, unsigned end,unsigned *total_nodes,
	std::vector<primitive_info_t>* ordered_primitives, const std::vector<float>& shape_data,int depth);
unsigned flatten_bvh_tree(const bvh_accelerator_t *bvh, bvh_build_node_t *node, unsigned* offset);

void build_bvh_accelerator(bvh_accelerator_t *bvh, const std::vector<primitive_info_t>& primitive_info_array,
	const std::vector<float>& shape_data)
{
		//TODO
	bvh->max_primitives_in_node = 2;
	std::vector<triangle_t> triangles;
	triangles.reserve(primitive_info_array.size());
	bvh->primitive_array = primitive_info_array;

	triangle_t tri;

	for(size_t i = 0;i < primitive_info_array.size();++i)
	{
		load_triangle(&tri,primitive_info_array[i].shape_info,shape_data);
		triangles.push_back(tri);
	}
	if(triangles.size() == 0)
	{
		bvh->nodes = 0;
		return;
	}
	std::vector<bvh_primitive_info_t> build_data;
	build_data.reserve(triangles.size());
	for(size_t i = 0;i < triangles.size(); ++i)
	{
		bbox_t bbox = world_bound(primitive_info_array[i].shape_info,shape_data);
		build_data.push_back(bvh_primitive_info_t(i, bbox));
	}
	std::vector<primitive_info_t> ordered_primitives;
	ordered_primitives.reserve(primitive_info_array.size());
	unsigned total_nodes = 0;
	bvh_build_node_t *root = recursive_build_bvh(bvh,build_data,0,primitive_info_array.size(), &total_nodes,&ordered_primitives,shape_data, 0);
	bvh->primitive_array.swap(ordered_primitives);
	bvh->nodes = new linear_bvh_node_t[total_nodes];
	unsigned offset = 0;
	flatten_bvh_tree(bvh,root, &offset);
	assert(offset == total_nodes);
	bvh->linear_node_num = total_nodes;
}
bvh_build_node_t *recursive_build_bvh(bvh_accelerator_t *bvh, std::vector<bvh_primitive_info_t>& build_data,unsigned start, unsigned end,unsigned *total_nodes,
	std::vector<primitive_info_t>* ordered_primitives, const std::vector<float>& shape_data, int depth)
{
	(*total_nodes) ++;
	bvh_build_node_t *node = new bvh_build_node_t;
	bbox_t bbox;
	bbox_init(&bbox);
	for(unsigned i = start; i < end; ++i)
	{
		bbox_union_with_bbox(&bbox, &build_data[i].bound);
	}
	unsigned n_primitives = end - start;
	if (n_primitives == 1 || depth >= MAX_BVH_DEPTH)
	{
		unsigned first_primitive_offset = ordered_primitives->size();
		for(unsigned i = start; i < end; ++i)
		{
			unsigned primitive_number = build_data[i].primitive_index;
			ordered_primitives->push_back(bvh->primitive_array[primitive_number]);
		}
		bvh_build_node_init_leaf(node,first_primitive_offset,n_primitives,bbox);
	}
	else
	{
		bbox_t centroid_bounds;
		bbox_init(&centroid_bounds);
		for(unsigned i = start; i < end; ++i)
			bbox_union_with_point(&centroid_bounds, &build_data[i].centroid);
		int dim = bbox_max_extent(&centroid_bounds);

		unsigned mid = (start + end) /2;
		if(centroid_bounds.pmax[dim] == centroid_bounds.pmin[dim])
		{
			unsigned first_primitive_offset = ordered_primitives->size();
			for(unsigned i = start; i < end; ++i)
			{
				unsigned primitive_number = build_data[i].primitive_index;
				ordered_primitives->push_back(bvh->primitive_array[primitive_number]);
			}
			bvh_build_node_init_leaf(node,first_primitive_offset,n_primitives,bbox);
			return node;
		}
		//sah
		if (n_primitives <= 4)
		{
			mid = (start + end)/2;
			std::nth_element(&build_data[start], &build_data[mid], &build_data[end-1]+1, 
				[=](const bvh_primitive_info_t& a, const bvh_primitive_info_t& b) ->bool {
					return a.centroid[dim] < b.centroid[dim];
			}
				);
		}
		else
		{
			const unsigned int n_buckets = 12;
			struct bucket_info_t{
				bucket_info_t(){
					count = 0;
					bbox_init(&bounds);
				}
				int count;
				bbox_t bounds;
			};
			bucket_info_t buckets[n_buckets];
			for(unsigned i = start; i < end; ++i)
			{
				unsigned int b = static_cast< unsigned > (n_buckets * (
					(build_data[i].centroid[dim] - centroid_bounds.pmin[dim])/(centroid_bounds.pmax[dim] - centroid_bounds.pmin[dim])
					));
				if (b == n_buckets) b = n_buckets - 1;
				assert(b >= 0 && b < n_buckets);
				buckets[b].count ++;
				bbox_union_with_bbox(&buckets[b].bounds, &build_data[i].bound);
			}

				float cost[n_buckets - 1];
				for(int i = 0; i < n_buckets - 1; ++i)
				{
					bbox_t b0,b1;
					bbox_init(&b0);
					bbox_init(&b1);
					int count0 = 0,count1 = 0;
					for (int j = 0;j <= i; ++j)
					{
						bbox_union_with_bbox(&b0, &buckets[j].bounds);	
						count0 += buckets[j].count;
					}
					for(int j= i+1;j < n_buckets; ++j)
					{
						bbox_union_with_bbox(&b1, &buckets[j].bounds);
						count1 += buckets[j].count;
					}
					cost[i] = .125f +(count0 * bbox_surface_area(&b0) + count1*bbox_surface_area(&b1))/bbox_surface_area(&bbox);
				}
				float min_cost = cost[0];
				unsigned min_cost_split = 0;
				for(int i = 1; i < n_buckets - 1; ++i)
				{
					if(cost[i] < min_cost){
						min_cost = cost[i];
						min_cost_split = i;
					}
				}
				if(n_primitives > bvh->max_primitives_in_node ||
					min_cost < n_primitives)
				{
					const bvh_primitive_info_t *pmid = std::partition(&build_data[start],
						&build_data[end-1]+1,
						[=](const bvh_primitive_info_t &p)->bool{
							unsigned int b = static_cast<unsigned int>(n_buckets * ((p.centroid[dim] - centroid_bounds.pmin[dim])/(centroid_bounds.pmax[dim]-centroid_bounds.pmin[dim])));
							if(b == n_buckets) b = n_buckets - 1;
							return b <= min_cost_split;
					}
					);
					mid = pmid - &build_data[0];
				}
				else
				{
					unsigned first_primitive_offset = ordered_primitives->size();
					for(unsigned i = start; i < end; ++i)
					{
						unsigned primitive_number = build_data[i].primitive_index;
						ordered_primitives->push_back(bvh->primitive_array[primitive_number]);
					}
					bvh_build_node_init_leaf(node,first_primitive_offset,n_primitives,bbox);
				}
		}
		bvh_build_node_init_interior(node, dim, 
			recursive_build_bvh(bvh, build_data, start,mid,total_nodes, ordered_primitives,shape_data,depth+1),
			recursive_build_bvh(bvh, build_data, mid,end,total_nodes, ordered_primitives,shape_data,depth+1));
	}
	return node;
}

unsigned flatten_bvh_tree(const bvh_accelerator_t *bvh, bvh_build_node_t *node, unsigned *offset)
{
	linear_bvh_node_t *linear_node = &bvh->nodes[*offset];
	linear_node->bounds = node->bounds;
	unsigned my_offset = (*offset) ++;
	if(node->n_primitives > 0)
	{
		assert(!node->children[0] && !node->children[1]);
		linear_node->primitive_offset                     = node->first_primitive_offset;
		linear_node->n_primitives = node->n_primitives;
	}
	else
	{
		linear_node->axis = node->split_axis;
		linear_node->n_primitives = node->n_primitives;
		flatten_bvh_tree(bvh,node->children[0], offset);
		linear_node->second_child_offset = flatten_bvh_tree(bvh,node->children[1], offset);
	}
	return my_offset;
}
void destroy_bvh_accelerator(bvh_accelerator_t **bvh)
{
	delete [] (*bvh)->nodes;
	delete (*bvh);
	*bvh = NULL;
}